# -*- coding: utf-8 -*-
"""
Created on Mon Sep  5 11:17:02 2022

@author: Runx

Function: a fine resolution spatial-temporal 
          rainfall generator
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy import stats
import math
import settings

# ws and ls must be an integral multiple of the resolution.
ws = settings.ws #storm envelope width, (meter)
ls = settings.ls #storm envelope length, (meter)
resolution = settings.resolution
num= settings.num

def num_gen():
    '''
    Generate the number of storm cells.
    降雨单元的数量，第个降雨单元被设置为 disc or ellipse(s1, s2).
    '''
    lambd = stats.norm.rvs(loc=0.0099, scale=0.0032, size=1) 
    if lambd <= 0:
        lambd = -lambd
    #calculate lambda，λ，Number of cells per square kilometre，km^(-2)
    # Note that the unit is km for ws and ls.
    num =  stats.poisson.rvs(lambd*ws*ls/10**6)
    if num == 0:
        num = 1
    return(num)
    
def loc_gen(num=10):
    '''
    location of storm cells, sampling from U([0, l_s]) and 
    U([0,w_s]), yields the coordinates x_c and y_c
    
    '''
    grid_n = ws//resolution
    grid_m = ls//resolution
    x = stats.randint.rvs(0, grid_m, size=num)
    y = stats.randint.rvs(0, grid_n, size=num)
    coordinates = {"x":x, "y":y}
    cell_loc = pd.DataFrame(coordinates)
    return(cell_loc)

def duration_gen():
    '''
    Storm duration using the Gamma distribution, (day)
    暴雨持续时长, Gamma分布
    '''
    a = 2.44 
    b = 17.92 
    d = stats.gamma.rvs(a, scale=b)
    #这里有个一个问题，即McRobie的结果里b=17.92，貌似应该是sd=17.92，验证
    return(d)

def s1_gen(num):
    '''
    Cell spread in direction of motion.
    用Log normal分布
    
    对数正态分布中，lg_N(s, loc, scale), 期中mu
    '''
    #global num 
    mu = 0.700
    sigma = 1.060
    s1 = stats.lognorm.rvs(s=sigma, loc=mu, scale=math.exp(mu), size=num)
    s1 = s1.astype("float32")
    return(s1)

def s2_gen(num):
    '''
    Cell spread in direction perpendicular to motion.
    用Log normal分布
    
    对数正态分布中，lg_N(s, loc, scale), 期中mu
    '''
    #global num 
    mu = 0.704
    sigma = 1.001
    s2 = stats.lognorm.rvs(s=sigma, loc=mu, scale=math.exp(mu), size=num)
    s2 = s2.astype("float32")
    return(s2)

def r_gen(num):
    '''
    Maximum intensity in cell (rmax)
    using Generalised Pareto distribution
    使用genpareto函数，a generalized Pareto continuous random variable
    
    $y=f(x \mid k, \sigma, \theta)=\left(\frac{1}{\sigma}\right)\left(1+
         k \frac{(x-\theta)}{\sigma}\right)^{-1-\frac{1}{k}}$
    k是scipy中的c，as the shape parameter
    theta是scipy中的loc
    '''
    k = 0.586
    sigma = 1.040
    theta = 0.2
    rmax = stats.genpareto.rvs(c=k, loc=theta, scale=sigma, size=num)
    rmax = rmax.astype("float32")
    return(rmax)

def velocity_gen():
    '''
    Velocity magnitude, (km/hr)
    
    Suppose a normally distributed random variable X has mean mu and \
    standard deviation sigma. Then Y = exp(X) is lognormally distributed \
    with s = sigma and scale = exp(mu).
    '''
    mean = 78.30
    sd = 23.64
    
    mu = math.log(mean**2/math.sqrt(sd**2+mean**2))
    sigma = math.sqrt(math.log(sd**2/mean**2+1))
    
    velocity = stats.lognorm.rvs(s=sigma, loc=mu, scale=math.exp(mu))
    return(velocity)

def direction_gen():
    '''
    Velocity direction, (rads)
    '''
    mean = 0.508
    sd = 1.164
    mu = math.log(mean**2/math.sqrt(sd**2+mean**2))
    sigma = math.sqrt(math.log(sd**2/mean**2+1))
    direction = stats.lognorm(s=sigma, loc=mu, scale=math.exp(mu))
    return(direction)

def max_num_gen():
    '''
    To evaluate the approximate max_num of storm cells in the system. 
    '''
    tmp=[]
    for i in range(10000):
        tmp.append(num_gen())
    tmp = np.array(tmp)
    return(tmp.max())    